package state

import (
	"osiris/dao"
	"osiris/dto"
	"osiris/logger"
	"osiris/model"
	"osiris/mpt"
	"osiris/ocrypto/ecdsa"
	"sync"
)

// AccountStateTrie 账户状态树(key: 账户地址对应的字节数组，value: dto.StateLeaf)
type AccountStateTrie struct {
	BaseStateTrie

	// nonceMap 记录每一个叶子节点的nonce的map
	nonceMap map[string]int64
}

func (trie *AccountStateTrie) InitState() error {
	trie.stateMutex = new(sync.RWMutex)
	trie.stateTrie = mpt.NewTrie()
	trie.nonceMap = make(map[string]int64, 0)

	accountModels, err := dao.GetAllAccount()
	if err != nil {
		logger.Error(map[string]interface{}{"[state] [Init State Trie] dao.GetAllAccount()": err})
		return err
	}

	if len(accountModels) == 0 {
		logger.Info(map[string]interface{}{"[state] [Init State Trie] ": "the length of accountModels in DB is 0"})
		return nil
	}

	var errInsideFor error
	errInsideFor = nil
	for _, accountModel := range accountModels {

		leaf := dto.StateLeaf{
			Asset:     accountModel.Asset,
			CoinSince: accountModel.CoinSince,
			Nonce:     accountModel.Nonce,
			MaxNonce:  accountModel.MaxNonce,
		}

		key := ecdsa.StrToAccountBytes(accountModel.Addr)
		err := trie.commitState(key, leaf, false)
		if err != nil {
			errInsideFor = err
			logger.Error(map[string]interface{}{"[state] [Init State] write state": err})
		}
	}

	return errInsideFor
}

func (trie *AccountStateTrie) ResetState() {
	trie.stateMutex.Lock()
	trie.stateTrie = mpt.NewTrie()
	trie.nonceMap = make(map[string]int64, 0)
	trie.stateMutex.Unlock()
}

// commitState 变更某一个账户的状态，包括状态树的叶子、nonceMap（对状态树的修改操作都要收敛到这个方法)
//
//	@param addrStr 账户地址
//	@param leaf 状态树叶子
//	@param writeDB 是否需要写入DB
//	@return error
func (trie *AccountStateTrie) commitState(key []byte, leaf dto.StateLeaf, writeDB bool) error {
	leafBytes, err := dto.ToBytes(leaf)
	if err != nil {
		logger.Error(map[string]interface{}{"[state] [change State] from StateLeaf to bytes": err})
		return err
	}

	addrStr := ecdsa.BytesToAccountStr(key)
	if writeDB {
		go trie.storeAccountState(addrStr, leaf)
	}

	trie.stateTrie.Put(key, leafBytes)
	trie.nonceMap[addrStr] = leaf.Nonce

	return nil
}

// storeAccountState 将账户状态变动写入MySQL
//
//	@param addrStr
//	@param leaf
//	@return error
func (trie *AccountStateTrie) storeAccountState(addrStr string, leaf dto.StateLeaf) error {
	//先尝试更新已有数据
	rowAffected, err1 := dao.UpdateAccount(addrStr, leaf.Asset, leaf.CoinSince, leaf.Nonce)
	if err1 != nil {
		logger.Error(map[string]interface{}{"[state] [store State] dao.UpdateAccount()": err1})
		return err1
	}

	//是新数据的话增添数据
	if rowAffected == 0 {
		err2 := dao.AddAccountByModel(model.AccountModel{
			Addr:      addrStr,
			Asset:     leaf.Asset,
			CoinSince: leaf.CoinSince,
			Nonce:     leaf.Nonce,
			MaxNonce:  leaf.MaxNonce,
		})

		if err2 != nil {
			logger.Error(map[string]interface{}{"[state] [store State] dao.AddAccountByModel()": err2})
		}
		return err2
	}

	return nil
}

// GetNonce 查询某个账户当前的随机数
//
//	@param addrStr 账户地址字符串
//	@return int64 当前随机数
//	@return bool 账户是否存在
func (trie *AccountStateTrie) getNonce(addrStr string) (int64, bool) {
	trie.stateMutex.RLock()
	nonce, exist := trie.nonceMap[addrStr]
	trie.stateMutex.RUnlock()
	if !exist {
		return InvalidNonce, false
	}

	return nonce, true
}

// IsNonceValid 根据账户nonce判断交易是否过期
//
//  @receiver trie 
//  @param addrStr 账户地址
//  @param nonce 交易的nonce
//  @return TxStatus 交易状态
func (trie *AccountStateTrie) IsNonceValid(addrStr string, nonce int64) TxStatus {
	//账户不存在算不合法
	currentNonce, accountExist := trie.getNonce(addrStr)
	if !accountExist {
		return TxInvalid
	}

	if currentNonce == nonce {
		return TxPending
	}

	if currentNonce < nonce {
		return TxQueued
	}

	return TxInvalid
}
